!------------------------------------------------------------------------------
! TJU/Department of Mechanics, Fluid Mechanics, Code START
!------------------------------------------------------------------------------
!> @Author: Jianxin Liu shookware@tju.edu.cn
!> @Date: 2024-08-13 21:38:30
!> @LastEditTime: 2024-08-13 21:38:32
!> @LastEditors: Jianxin Liu shookware@tju.edu.cn
!> @Description:
!> @FilePath: /START_NS_3D/src/splitter/module_splitter_steger_warming.f90
!> @Copyright (c) 2024 by Jianxin Liu email: shookware@tju.edu.cn, All Rights Reserved.
!------------------------------------------------------------------------------
submodule(mod_splitter) submod_splitter_steger_warming
   implicit none

contains

   module subroutine SplitStegerWarming(this, q_in, q_l, q_r, chi_x, chi_y, chi_z, chi_face)
      class(type_splitter_steger_warming), intent(inout) :: this
      class(type_vector_base), intent(in) :: q_in(:)
      class(type_vector_base), intent(inout) :: q_l(:), q_r(:)
      class(type_vector_base), intent(in) :: chi_x, chi_y, chi_z, chi_face

      type(type_self_vector) :: uhat, c
      type(type_self_vector) :: lambda1_p, lambda2_p, lambda3_p
      type(type_self_vector) :: lambda1_m, lambda2_m, lambda3_m
      type(type_self_vector) :: eta_p(2:3), eta_m(2:3)
      type(type_self_vector) :: work1, work2, work3

      real, parameter :: eps = 1.0d-6

      integer :: n, l

      n = size(q_in)

      call uhat%duplicate(q_in(1))
      call c%duplicate(q_in(1))

      call lambda1_p%duplicate(q_in(1))
      call lambda2_p%duplicate(q_in(1))
      call lambda3_p%duplicate(q_in(1))
      call lambda1_m%duplicate(q_in(1))
      call lambda2_m%duplicate(q_in(1))
      call lambda3_m%duplicate(q_in(1))
      do l = 2, 3
         call eta_p(l)%duplicate(q_in(1))
         call eta_m(l)%duplicate(q_in(1))
      end do
      call work1%duplicate(q_in(1))
      call work2%duplicate(q_in(1))
      call work3%duplicate(q_in(1))

      associate (gamma => this%eos%ref%gamma)
         call uhat%pointwise_mult(chi_x, q_in(2))
         call work1%pointwise_mult(chi_y, q_in(3))
         call work2%pointwise_mult(chi_z, q_in(4))
         call uhat%axpbypcz(1.0d0, 1.0d0, 1.0d0, work1, work2) !<单位逆变速度

         !! 完全气体模型 否则需要计算T*cp/(cv*gamma) cp和cv都为有量纲的值(也就是gamma/gamma_ref), T为无量纲的量
         call c%copy(q_in(5))
         call c%sqrtabs() !< sqrt(T)
         call c%scale(this%eos%ref%flg_uref) !< 声速
      end associate
      !! lambda1, lambda2, lambda3
      associate (lambda1 => lambda1_p, lambda2 => lambda2_p, lambda3 => lambda3_p)
         call lambda2%pointwise_mult(uhat, chi_face) !< lambda2 = uhat*|\grad chi|/J
         call lambda1%copy(uhat)
         call lambda3%copy(uhat)
         call lambda1%axpy(-1.0d0, c)
         call lambda1%pointwise_mult(lambda1, chi_face)
         call lambda3%axpy(1.0d0, c)
         call lambda3%pointwise_mult(lambda3, chi_face)

         !!lambda-
         !> lambda1_m
         call lambda1_m%pointwise_mult(lambda1, lambda1)
         call lambda1_m%shift(eps**2)
         call lambda1_m%sqrtabs()
         call lambda1_m%aypx(-1.0d0, lambda1)
         call lambda1_m%scale(0.5d0)

         !> lambda2_m
         call lambda2_m%pointwise_mult(lambda2, lambda2)
         call lambda2_m%shift(eps**2)
         call lambda2_m%sqrtabs()
         call lambda2_m%aypx(-1.0d0, lambda2)
         call lambda2_m%scale(0.5d0)

         !> lambda3_m
         call lambda3_m%pointwise_mult(lambda3, lambda3)
         call lambda3_m%shift(eps**2)
         call lambda3_m%sqrtabs()!
         call lambda3_m%aypx(-1.0d0, lambda3)
         call lambda3_m%scale(0.5d0)
      end associate

      call lambda1_p%axpy(-1.0d0, lambda1_m) !< lambda{i}_p = lambda{i} - lambda{i}_m
      call lambda2_p%axpy(-1.0d0, lambda2_m)
      call lambda3_p%axpy(-1.0d0, lambda3_m)

      !> eta
      #: for l, s in list(zip(["l", "r"], ["p", "m"]))
      call q_${l}$ (1)%copy(lambda2_${s}$)
      call q_${l}$ (1)%scale(2.0d0*this%eos%ref%gamma)
      call q_${l}$ (1)%axpy(-2.0d0, lambda2_${s}$)
      call q_${l}$ (1)%axpy(1.0d0, lambda1_${s}$)
      call q_${l}$ (1)%axpy(1.0d0, lambda3_${s}$)
      call eta_${s}$ (2)%copy(lambda3_${s}$)
      call eta_${s}$ (2)%axpy(-1.0d0, lambda1_${s}$)
      call eta_${s}$ (3)%copy(lambda3_${s}$)
      call eta_${s}$ (3)%axpy(1.0d0, lambda1_${s}$)
      #: endfor

      #: for l, s in list(zip(["l", "r"], ["p", "m"]))
      call work1%pointwise_mult(eta_${s}$ (2), c)

      #: for i, x in list(zip(range(2, 5), ["x", "y", "z"]))
      call work2%pointwise_mult(q_${l}$ (1), q_in(${i}$))
      call q_${l}$ (${i}$)%pointwise_mult(work1, chi_${x}$)
      call q_${l}$ (${i}$)%axpy(1.0d0, work2)
      #: endfor
      call q_${l}$ (5)%pointwise_mult(c, c)
      call q_${l}$ (5)%pointwise_mult(q_${l}$ (5), eta_${s}$ (3))
      call q_${l}$ (5)%scale(1.0d0/(this%eos%ref%gamma - 1.0d0))
      call work1%pointwise_mult(work1, uhat)
      call q_${l}$ (5)%axpy(1.0d0, work1)
      call work1%pointwise_mult(q_in(2), q_in(2))
      call work2%pointwise_mult(q_in(3), q_in(3))
      call work3%pointwise_mult(q_in(4), q_in(4))
      call work3%axpbypcz(0.5d0, 0.5d0, 0.5d0, work1, work2)
      call work3%pointwise_mult(q_${l}$ (1), work3)
      call q_${l}$ (5)%axpy(1.0d0, work3)
      do l = 1, 5
         call q_${l}$ (l)%scale(1.0d0/(2.0d0*this%eos%ref%gamma))
         call q_${l}$ (l)%pointwise_mult(q_${l}$ (l), q_in(1))
      end do
      #:endfor

      call uhat%destroy()
      call c%destroy()
      call lambda1_p%destroy()
      call lambda2_p%destroy()
      call lambda3_p%destroy()
      call lambda1_m%destroy()
      call lambda2_m%destroy()
      call lambda3_m%destroy()
      do l = 2, 3
         call eta_p(l)%destroy()
         call eta_m(l)%destroy()
      end do
      call work1%destroy()
      call work2%destroy()
      call work3%destroy()

      ! block
      !    print*, petsc_get_rank(), vector_equal(q_in(1), q_in(2)), &
      !    vector_equal(q_in(1), q_in(3)), &
      !    vector_equal(q_in(1), q_in(4)), &
      !    vector_equal(q_in(1), q_in(5)), &
      !    vector_equal(q_in(2), q_in(3)), &
      !    vector_equal(q_in(2), q_in(4)), &
      !    vector_equal(q_in(2), q_in(5)), &
      !    vector_equal(q_in(3), q_in(4)), &
      !    vector_equal(q_in(3), q_in(5)), &
      !    vector_equal(q_in(4), q_in(5))
      ! endblock

      end subroutine SplitStegerWarming

      subroutine SplitStegerWarming2(this, q_in, q_l, q_r, chi_x, chi_y, chi_z, chi_face)
         implicit none
         class(type_splitter_steger_warming), intent(inout) :: this
         type(type_self_vector), intent(in) :: q_in(:)
         type(type_self_vector), intent(inout) :: q_l(:), q_r(:)
         type(type_self_vector), intent(in) :: chi_x, chi_y, chi_z, chi_face

         type(type_array) :: v_in(size(q_in)), v_l(size(q_in)), v_r(size(q_in))
         type(type_array) :: k_x, k_y, k_z, k_face
         real(dp) :: vv_in(size(q_in)), vv_l(size(q_in)), vv_r(size(q_in))

         integer :: l, lsta, lend, i

         do l = 1, size(q_in)
            call q_in(l)%get_array_read(v_in(l)%p)
            call q_l(l)%get_array(v_l(l)%p)
            call q_r(l)%get_array(v_r(l)%p)
         end do

         call chi_x%get_array_read(k_x%p)
         call chi_y%get_array_read(k_y%p)
         call chi_z%get_array_read(k_z%p)
         call chi_face%get_array_read(k_face%p)

         do concurrent(l=lbound(v_in(1)%p, dim=1):ubound(v_in(1)%p, dim=1))
            do concurrent(i=1:size(q_in))
               vv_in(i) = v_in(i)%p(l)
            end do
            call SplitPoint(this, vv_in, vv_l, vv_r, k_x%p(l), k_y%p(l), k_z%p(l), k_face%p(l))
            do concurrent(i=1:size(q_in))
               v_l(i)%p(l) = vv_l(i)
               v_r(i)%p(l) = vv_r(i)
            end do
         end do

         call chi_x%restore_array_read(k_x%p)
         call chi_y%restore_array_read(k_y%p)
         call chi_z%restore_array_read(k_z%p)
         call chi_face%restore_array_read(k_face%p)
         do l = 1, size(q_in)
            call q_in(l)%restore_array_read(v_in(l)%p)
            call q_l(l)%restore_array(v_l(l)%p)
            call q_r(l)%restore_array(v_r(l)%p)
         end do
      end subroutine SplitStegerWarming2

      pure subroutine SplitPoint(this, v_in, v_l, v_r, k_x, k_y, k_z, k_face)
         implicit none
         class(type_splitter_steger_warming), intent(in) :: this
         real(dp), intent(in) :: v_in(:)
         real(dp), intent(out) :: v_l(:), v_r(:)
         real(dp), intent(in) :: k_x, k_y, k_z, k_face

         real(dp) :: rho, u, v, w, T
         real(dp) :: ubar, c
         real(dp) :: lambda(3, 2)
         real(dp) :: eta(3, 2)
         real(dp) :: tmp1, tmp2, tmp3
         integer :: i

         real(dp), parameter :: eps = 1.0d-6

         rho = v_in(1)
         u = v_in(2); v = v_in(3); w = v_in(4)
         T = v_in(5)

         ubar = k_x*u + k_y*v + k_z*w
         c = sqrt(T)*this%eos%ref%flg_uref

         tmp1 = ubar - c; tmp2 = ubar; tmp3 = ubar + c
         lambda(1, 1) = 0.5d0*(tmp1 + sqrt((tmp1 + eps)**2))
         lambda(2, 1) = 0.5d0*(tmp2 + sqrt((tmp2 + eps)**2))
         lambda(3, 1) = 0.5d0*(tmp3 + sqrt((tmp3 + eps)**2))
         lambda(1, 2) = tmp1 - lambda(1, 1)
         lambda(2, 2) = tmp2 - lambda(2, 1)
         lambda(3, 2) = tmp3 - lambda(3, 1)

         lambda = lambda*k_face

         do i = 1, 2
            eta(1, i) = 2.0d0*this%eos%ref%gamma*lambda(2, i) + lambda(3, i) - 2.0d0*lambda(2, i) + lambda(1, i)
            eta(2, i) = lambda(3, i) - lambda(1, i)
            eta(3, i) = lambda(3, i) + lambda(1, i)
         end do

         tmp1 = eta(2, 1)*c
         v_l(1) = eta(1, 1)
         v_l(2) = eta(1, 1)*u + tmp1*k_x
         v_l(3) = eta(1, 1)*v + tmp1*k_y
         v_l(4) = eta(1, 1)*w + tmp1*k_z
         v_l(5) = eta(1, 1)*0.5d0*(u**2 + v**2 + w**2) + tmp1*ubar + eta(3, 1)*c**2/(this%eos%ref%gamma - 1.0d0)

         tmp1 = eta(2, 2)*c
         v_r(1) = eta(1, 2)
         v_r(2) = eta(1, 2)*u + tmp1*k_x
         v_r(3) = eta(1, 2)*v + tmp1*k_y
         v_r(4) = eta(1, 2)*w + tmp1*k_z
         v_r(5) = eta(1, 2)*0.5d0*(u**2 + v**2 + w**2) + tmp1*ubar + eta(3, 2)*c**2/(this%eos%ref%gamma - 1.0d0)

         v_l = v_l*rho/(2.0d0*this%eos%ref%gamma)
         v_r = v_r*rho/(2.0d0*this%eos%ref%gamma)

      end subroutine SplitPoint

   end submodule submod_splitter_steger_warming
